import sys
import Detect
import cv2 as cv
import math
import time
import numpy as np
import torch
from matplotlib import pyplot as plt
from PyQt5 import QtGui
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QMessageBox
from PyQt5.QtWidgets import QMainWindow, QLabel
from PyQt5.QtGui import QPixmap, QImage, QPainter
from Detect import Ui_MainWindow
from PyQt5.QtCore import Qt


# opencv读取图片转换Qt图片
def Opencv2QImage(cvimg):
    height, width, depth = cvimg.shape
    cvimg = cv.cvtColor(cvimg, cv.COLOR_BGR2RGB)
    cvimg = QImage(cvimg.data, width, height, width * depth, QImage.Format_RGB888)
    return cvimg


# qt图片转换opencv读取图片
def QPixmap2Opencv(pixmap):
    qimg = pixmap.toImage()
    temp_shape = (qimg.height(), qimg.bytesPerLine() * 8 // qimg.depth())
    temp_shape += (4,)
    ptr = qimg.bits()
    ptr.setsize(qimg.byteCount())
    result = np.array(ptr, dtype=np.uint8).reshape(temp_shape)
    result = result[..., :3]
    return result


# 链接网络摄像头
def get_DroidCam_url(ip="192.168.137.11", port=4747, res='1080p'):
    res_dict = {
        '240p': '320x240',
        '360p': '480*360',
        '480p': '640x480',
        '720p': '1280x720',
        '1080p': '1920x1080',
    }
    url = f'http://{ip}:{port}/mjpegfeed?{res_dict[res]}'
    return url


# 模型加载
detecter = torch.hub.load('../yolov5', 'custom', path='../yolov5/runs/train/exp12/weights/best.pt', source='local')


class MyMainWindow(QMainWindow, Ui_MainWindow):  # 继承 QMainWindow类和 Ui_MainWindow界面类
    def __init__(self, parent=None):
        super(MyMainWindow, self).__init__(parent)  # 初始化父类
        self.setupUi(self)  # 继承 Ui_MainWindow 界面类
        self.cap = None  # 摄像头

    def OpenPicture(self):
        # 清空
        self.Viedo_Label.clear()
        # 打开文件资源管理器选择图片
        imagePath, _ = QFileDialog.getOpenFileName(None, "Open File", "img", "JPEG Files(*.jpg);;Images (*.jpg *.png)")
        if not imagePath:
            QMessageBox.warning(self, 'Error', '请选择好一张图片.')
            return
        img = cv.imread(imagePath)  # 通过Opencv读入一张图片
        print(int(img.shape[0]), int(img.shape[1]))
        if  int(img.shape[0]) > 640 or  int(img.shape[1]) > 480:# 如果图片过大，进行裁剪
            img = cv.resize(img, (640, int(640 * img.shape[0] / img.shape[1])), interpolation=cv.INTER_AREA)
        self.Detect(img)

    def OpenVideo(self):
        # 清空
        self.Viedo_Label.clear()
        # 打开文件资源管理器选择视频
        videoPath, _ = QFileDialog.getOpenFileName(None, 'Open File', 'img', 'Video Files (*.mp4 *.avi)')
        if not videoPath:
            QMessageBox.warning(self, 'Error', '请选择视频.')
            return
        # 将视频读入cap
        self.cap = cv.VideoCapture(videoPath)

        while True:
            ret, img = self.cap.read()
            if ret is False:
                del ret
                del img
                break

            # 缩放视频大小
            img = cv.resize(img, (640, int(640 * img.shape[0] / img.shape[1])), interpolation=cv.INTER_AREA)
            # 图像处理
            self.Detect(img)
            # 释放图片，防止内存占用
            del ret
            del img
            # 关闭窗体
            key = cv.waitKey(50)
            if key == 27:  # esc
                break
        # 释放摄像头
        self.cap.release()
        self.Viedo_Label.clear()
        return

    def OpenCapture(self):
        self.Viedo_Label.clear()
        # 变量定义
        fps = 0                     # 帧数
        frames_counter = 0          # 帧数计时器
        start_time = time.time()    # 开始时间

        ip = self.Capture_IP.toPlainText()
        url = get_DroidCam_url(ip, 4747, '240p')
        self.cap = cv.VideoCapture(url)

        while self.cap.isOpened():  # 有视频输入
            ret, img = self.cap.read()
            if ret is False:
                del ret
                del img
                break
            # 帧数测量
            if ret == 1:
                frames_counter += 1
            elapsed_time = time.time() - start_time
            if elapsed_time >= 1.0:         # 1秒
                fps = frames_counter
                frames_counter = 0
                start_time = time.time()
            # 输出fps
            cv.putText(img, f"fps:{fps}", (250, 30), cv.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1)
            # 图像处理
            self.Detect(img)
            # 释放图片，防止内存占用
            del ret
            del img
            # 关闭窗体
            key = cv.waitKey(50)
            if key == 27:  # esc
                break
        # 释放摄像头
        self.cap.release()
        return

    def StopCapture(self):
        self.cap.release()
        return

    def Detect(self, img):
        gray_img = img.copy()
        gray_img = cv.cvtColor(img, cv.COLOR_RGB2GRAY)
        counter = 0     # 图片中天牛个数
        # 加载模型
        longicorns = detecter(gray_img)
        for longicorn in longicorns.xyxy[0]:
            # 天牛个数增加
            counter += 1
            # 获取到当前图片真正的xyxy，原来的格式应该是tensor，需要转化为列表
            longicorn = longicorn.tolist()
            xStart = longicorn[0]
            yStart = longicorn[1]
            xEnd = longicorn[2]
            yEnd = longicorn[3]
            cv.rectangle(img, (int(xStart), int(yStart)), (int(xEnd), int(yEnd)), (0, 255, 0), 3)
            # 检测置信度
            conf = longicorn[4]
            # 检测类别
            cls = longicorn[5]

        # 输出天牛个数
        self.Longicorn_Count.setHtml(f"<font size = 6>{str(counter)}</font>")
        # 加载处理后图片
        Qimg = Opencv2QImage(img)
        pixmap = QtGui.QPixmap.fromImage(Qimg)
        self.Viedo_Label.setPixmap(pixmap)
        self.Viedo_Label.setAlignment(Qt.AlignCenter)  # 将图像居中显示
        self.Viedo_Label.show()


if __name__ == '__main__':
    app = QApplication(sys.argv)  # 在 QApplication 方法中使用，创建应用程序对象
    myWin = MyMainWindow()  # 实例化 MyMainWindow 类，创建主窗口
    myWin.show()  # 在桌面显示控件 myWin
    sys.exit(app.exec_())  # 结束进程，退出程序
